import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.metrics import confusion_matrix, classification_report
from sklearn.svm import SVC
import matplotlib.pyplot as plt
from scipy.stats import skew, kurtosis
from scipy.fftpack import fft
import seaborn as sns

# 设置中文字体和负号显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体 (SimHei) 字体
plt.rcParams['axes.unicode_minus'] = False    # 解决坐标轴负号显示问题


# 1. 数据导入
data = pd.read_excel('附件一（训练集）材料1.xlsx')
data = data.drop(index=0)  # 删除第一行

trainData = data


# 2. 提取特征
def extract_features(data):
    # 提取基础特征
    temp = data.iloc[:, 0].values  # 温度
    freq = data.iloc[:, 1].values  # 频率
    core_loss = data.iloc[:, 2].values  # 磁芯损耗

    # 提取磁通密度数据
    flux_density = data.iloc[:, 4:].values  # 第5到1029列

    # 统计特征
    mean_flux = np.mean(flux_density, axis=1)  # 均值
    std_flux = np.std(flux_density, axis=1)  # 标准差
    max_flux = np.max(flux_density, axis=1)  # 最大值
    min_flux = np.min(flux_density, axis=1)  # 最小值
    peak_to_peak = max_flux - min_flux  # 峰峰值
    skewness_flux = skew(flux_density, axis=1)  # 偏度
    kurtosis_flux = kurtosis(flux_density, axis=1)  # 峰度

    # 傅里叶变换特征
    fft_flux = np.abs(fft(flux_density, axis=1))
    fft_mean = np.mean(fft_flux[:, 1:], axis=1)  # 排除直流分量后的均值
    fft_std = np.std(fft_flux[:, 1:], axis=1)  # 排除直流分量后的标准差

    # 组合所有特征
    features = np.column_stack([temp, freq, core_loss, mean_flux, std_flux, max_flux, min_flux,
                                peak_to_peak, skewness_flux, kurtosis_flux, fft_mean, fft_std])
    return features


X = extract_features(trainData)
y = pd.factorize(trainData.iloc[:, 3])[0]  # 将波形类型转换为数字

# 3. 划分训练集和验证集
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)

# 4. 模型训练（使用支持向量机SVM）
model = SVC(kernel='rbf', probability=True, random_state=42)
model.fit(X_train, y_train)

# 5. 验证模型
y_pred = model.predict(X_val)

# 5.1 输出分类报告
print('Classification Report:')
print(classification_report(y_val, y_pred))

# 5.2 绘制带有数字的混淆矩阵
cm = confusion_matrix(y_val, y_pred)
print('Confusion Matrix:')
print(cm)

# 使用 Seaborn 的 heatmap 绘制混淆矩阵，带有数字标注
plt.figure(figsize=(8, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=['正弦波', '三角波', '梯形波'], yticklabels=['正弦波', '三角波', '梯形波'])
plt.title('Confusion Matrix - Validation Set')
plt.ylabel('True label')
plt.xlabel('Predicted label')
plt.show()

# 6. 对附件二中的样本进行分类
testData = pd.read_excel('附件二（测试集）.xlsx')
X_test = extract_features(testData)

# 7. 进行预测
predictions = model.predict(X_test)

# 8. 显示预测结果
class_names = ['正弦波', '三角波', '梯形波']
predicted_classes = [class_names[p] for p in predictions]
print('Predicted Waveforms:')
print(predicted_classes)

# 9. 统计每种波形的数量
num_sine = np.sum(predictions == 0)
num_triangle = np.sum(predictions == 1)
num_trapezoidal = np.sum(predictions == 2)
print(f'正弦波数量: {num_sine}')
print(f'三角波数量: {num_triangle}')
print(f'梯形波数量: {num_trapezoidal}')
